import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;

import java.io.*;
import java.util.Map;

import static org.junit.Assert.*;

public class HuffTest {
    private Huff huff;
    private static final String TEST_INPUT_FILE = "/test_input.txt";
    private static final String TEST_COMPRESSED_FILE = "/test_compressed";
    private static final String TEST_UNCOMPRESSED_FILE = "/test_uncompressed.txt";

    @Rule
    public ExpectedException exception = ExpectedException.none();

    @Before
    public void setUp() {
        huff = new Huff();
        try {
            BufferedWriter writer = new BufferedWriter(new FileWriter(TEST_INPUT_FILE));
            writer.write("This is a test string with some repeated characters. " + "The more " +
                    "repetition, the better compression!");
            writer.close();
        } catch (IOException e) {
            fail("Could not create test file: " + e.getMessage());
        }
    }

    @After
    public void tearDown() {
        // clean up test files
        new File(TEST_INPUT_FILE).delete();
        new File(TEST_COMPRESSED_FILE).delete();
        new File(TEST_UNCOMPRESSED_FILE).delete();
    }

    @Test
    public void testBasicCompression() {
        int compressedSize = huff.write(TEST_INPUT_FILE, TEST_COMPRESSED_FILE, true);
        assertTrue("Compressed file should be created", new File(TEST_COMPRESSED_FILE).exists());
        assertTrue("Compressed size should be positive", compressedSize > 0);
    }

    @Test
    public void testBasicDecompression() throws IOException {
        huff.write(TEST_INPUT_FILE, TEST_COMPRESSED_FILE, true);

        int uncompressedChars = huff.uncompress(TEST_COMPRESSED_FILE, TEST_UNCOMPRESSED_FILE);
        assertTrue("Uncompressed file should be created",
                new File(TEST_UNCOMPRESSED_FILE).exists());

        String original = readFileToString(TEST_INPUT_FILE);
        String uncompressed = readFileToString(TEST_UNCOMPRESSED_FILE);
        assertEquals("Original and uncompressed files should match", original, uncompressed);
        assertEquals("Character count should match", original.length() * 8, uncompressedChars);
    }

    @Test
    public void testEmptyFile() throws IOException {
        writeStringToFile(TEST_INPUT_FILE, "");

        huff.write(TEST_INPUT_FILE, TEST_COMPRESSED_FILE, true);
        assertTrue("Should create compressed file even for empty input",
                new File(TEST_COMPRESSED_FILE).exists());

        huff.uncompress(TEST_COMPRESSED_FILE, TEST_UNCOMPRESSED_FILE);
        String content = readFileToString(TEST_UNCOMPRESSED_FILE);
        assertEquals("Decompressed empty file should be empty", "", content);
    }

    @Test
    public void testInvalidCompressedFile() throws IOException {
        writeStringToFile(TEST_COMPRESSED_FILE, "This is not a valid compressed file");

        int result = huff.uncompress(TEST_COMPRESSED_FILE, TEST_UNCOMPRESSED_FILE);

        assertEquals("Decompression should fail and return 0", 0, result);
    }

    @Test
    public void testForceCompression() throws IOException {
        String uniqueChars = "abcdefghijklmnopqrstuvwxyz";
        writeStringToFile(TEST_INPUT_FILE, uniqueChars);

        huff.write(TEST_INPUT_FILE, TEST_COMPRESSED_FILE, false);
        assertFalse("Compression should not occur without force flag when no space saved",
                new File(TEST_COMPRESSED_FILE).exists());

        huff.write(TEST_INPUT_FILE, TEST_COMPRESSED_FILE, true);
        assertTrue("Compression should occur with force flag",
                new File(TEST_COMPRESSED_FILE).exists());
    }

    @Test
    public void testEncodingTableGeneration() throws IOException {
        InputStream input = new ByteArrayInputStream("aabbc".getBytes());
        huff.makeHuffTree(input);
        input.close();

        Map<Integer, String> encodingMap = huff.makeTable();
        assertNotNull("Encoding map should not be null", encodingMap);
        assertTrue("Should have encoding for 'a'", encodingMap.containsKey((int) 'a'));
        assertTrue("Should have encoding for 'b'", encodingMap.containsKey((int) 'b'));
        assertTrue("Should have encoding for 'c'", encodingMap.containsKey((int) 'c'));
        assertTrue("Should have encoding for PSEUDO_EOF",
                encodingMap.containsKey(IHuffConstants.PSEUDO_EOF));
    }

    @Test
    public void testCharacterFrequencyCounts() throws IOException {
        String testString = "aaaaabbbcc";
        InputStream input = new ByteArrayInputStream(testString.getBytes());
        huff.makeHuffTree(input);
        input.close();

        Map<Integer, Integer> counts = huff.showCounts();
        assertEquals("Should count 5 'a' characters", Integer.valueOf(5), counts.get((int) 'a'));
        assertEquals("Should count 3 'b' characters", Integer.valueOf(3), counts.get((int) 'b'));
        assertEquals("Should count 2 'c' characters", Integer.valueOf(2), counts.get((int) 'c'));
    }

    @Test
    public void testCompressedSizeCalculation() throws IOException {
        String testString = "aaaaabbbcc";
        writeStringToFile(TEST_INPUT_FILE, testString);

        int compressedSize = huff.write(TEST_INPUT_FILE, TEST_COMPRESSED_FILE, true);

        File compressedFile = new File(TEST_COMPRESSED_FILE);
        assertTrue("Compressed file should exist", compressedFile.exists());
        assertTrue("Compressed size should be accurate",
                Math.abs(compressedSize / 8 - compressedFile.length()) <= 1);
    }

    @Test
    public void testLargeFile() throws IOException {
        StringBuilder content = new StringBuilder();
        for (int i = 0; i < 10240; i++) {
            content.append((char) ('a' + (i % 26)));
        }
        writeStringToFile(TEST_INPUT_FILE, content.toString());

        huff.write(TEST_INPUT_FILE, TEST_COMPRESSED_FILE, true);
        huff.uncompress(TEST_COMPRESSED_FILE, TEST_UNCOMPRESSED_FILE);

        String original = readFileToString(TEST_INPUT_FILE);
        String uncompressed = readFileToString(TEST_UNCOMPRESSED_FILE);
        assertEquals("Large file contents should match", original, uncompressed);
    }

    private void writeStringToFile(String filename, String content) throws IOException {
        BufferedWriter writer = new BufferedWriter(new FileWriter(filename));
        writer.write(content);
        writer.close();
    }

    private String readFileToString(String filename) throws IOException {
        StringBuilder content = new StringBuilder();
        BufferedReader reader = new BufferedReader(new FileReader(filename));
        String line;
        while ((line = reader.readLine()) != null) {
            content.append(line);
            if (reader.ready()) {
                content.append(System.lineSeparator());
            }
        }
        reader.close();
        return content.toString();
    }
}
